import torch
from torch import nn


class BatchNorm(nn.Module):
    def __init__(self, features:int, device):
        super(BatchNorm, self).__init__()
        self.features = features
        self.batch_norm = nn.BatchNorm1d(self.features).to(device=device)

    def forward(self, x:torch.Tensor) -> torch.Tensor:
        batch_size = x.shape[0]
        x = x.view(-1, self.features)
        x = self.batch_norm(x)
        x = x.view(batch_size, -1, self.features)
        return x